package cn.kbyue.demo.websocketdemo.controller;

import org.springframework.http.MediaType;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.io.IOException;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

@RestController
@RequestMapping(path = "sse")
public class WebSSEController {

    private static final Map<String, SseEmitter> SSE_CACHE = new ConcurrentHashMap<>();

    @GetMapping(path = "subscribe", produces = {MediaType.TEXT_EVENT_STREAM_VALUE})
    public SseEmitter push(String id) {
        // 超时时间设置为1小时, 0表示不过期。默认30秒，超过时间未完成会抛出异常：AsyncRequestTimeoutException
        SseEmitter sseEmitter = new SseEmitter(3600_000L);

        // 设置前端的重试时间为 1s
        // sseEmitter.send(SseEmitter.event().reconnectTime(1000).data("连接成功"));

        SSE_CACHE.put(id, sseEmitter);
        sseEmitter.onTimeout(() -> SSE_CACHE.remove(id));
        sseEmitter.onError(e -> {
            SSE_CACHE.remove(id);
            System.out.println(e.getMessage());
        });
        sseEmitter.onCompletion(() -> System.out.println("完成！！！"));
        return sseEmitter;
    }

    @GetMapping(path = "push")
    public String push(String id, String content) throws IOException {
        SseEmitter sseEmitter = SSE_CACHE.get(id);
        if (sseEmitter != null) {
            sseEmitter.send(content);
        }
        return "over";
    }

    @GetMapping(path = "over")
    public String over(String id) {
        SseEmitter sseEmitter = SSE_CACHE.get(id);
        if (sseEmitter != null) {
            sseEmitter.complete();
            SSE_CACHE.remove(id);
        }
        return "over";
    }

}
